from DriverInterface import DriverInterface
from math import *
from numpy import *
from numpy.linalg import *
from BonyLandmark import BonyLandmark
from Sensor import Sensor
import time
import fob
import time
import wx

import threading

##  Handles communication with the Flock of Birds system through the serial port.
class FoBDriver(DriverInterface):

    ##  The constructor of FoBDriver. You may only call this method if you inherit from DriverInterface.
    #   @param[in]  main        Main wx.App derived instance.
    #   @param[in]  numbirds    The number of birds connected to the XRC including the XRC itself.
    def __init__(self, main, numbirds = 10, fobrunning = True):
        if not isinstance(self, DriverInterface):
            raise TypeError("FoBDriver.__init__ requires DriverInterface.")

        ## Opens a workfile to write the raw data to.
        #self.fileWriter = open('workfile.txt', 'w')
        ##  Main wx.App derived instance.
        self.main = main
        self.rawoutputlist = []
        ##  Indicates whether the FoB system has started successfully.
        self.fobrunning = fobrunning
        ## Make the COM ports when fobrunning is set to true.
        if self.fobrunning:
            ##  A (key,value) mapping with comport numbers as key and corresponding FoB objects as value.
            self.flock = self.SetComports()
        ##  The list with five position measurements of each sensor connected to the corresponding bony landmarks.
        self.imSensors = [[] for i in range(len(self.main.bonyLandmarks) - len(self.main.GHJoints))]
        ##  The list with five position measurements of each bony landmark as measured with the stylus.
        #   Each entry in the list represents a bony landmark.
        self.stylus = [[] for i in range(len(self.main.bonyLandmarks) - len(self.main.GHJoints))]
        stylfile = main.prefReader.stylusFile
        try:
            stylusData = open(stylfile,'r').read()
        except:
            msg = "An error ocurred while opening %s." % stylfile
            dlg = wx.MessageDialog(self.main.gui, msg, 'Warning', wx.OK | wx.NO_DEFAULT | wx.ICON_EXCLAMATION)
            dlg.ShowModal()
        ##  The length of the stylus.
        self.__stylusLength = transpose(matrix("["+stylusData+"]"))/10

    ##  Reads the sensor:comport mapping from the config file and creates FoB objects.
    #   @return[dictionary] A (key,value) mapping with sensor ID's numbers as key and corresponding FoB objects as value.
    def SetComports(self):
        comports = self.main.config.ReadComports()
        ##  A (key,value) mapping with sensor ID's numbers as key and corresponding FoB objects as value.
        self.fob_dict = {}
        for gsi, comport in comports.items():
            comport = int(comport)
            gsi = int(gsi)
            if gsi != 1:
                #self.fob_dict[gsi] = self.init_FOB_parallel(comport)
        
                thread = threading.Thread(group = None, target = self.ThreadedFOBsetup, args = ([comport, gsi]))
                thread.start()
        
        while (len(threading.enumerate())>1):
            time.sleep(0.003)
            
        self.fob_dict[1] = self.init_XRC(10,int(comports['1']))

        return self.fob_dict

    ##  Creates an object for and starts the XRC.
    #   @param[in]  nr_of_birds The number of sensors connected to the XRC including the XRC itself.
    #   @param[in]  port        The comport where the XRC is plugged into.
    def init_XRC(self, nr_of_birds, port):
        f = fob.FOB()
        time.sleep(0.01)
        comport = port-1
        baudrate = 19200#19200
        # Default timeout in seconds.
        timeout = 5
        try: f.open(comport, baudrate, timeout)
        except:
            msg = "Could not open FoB connections: COM port %d does not exist on this computer.\nPlease reconfigure the FobDevice of the config.ini file." %comport
            dlg = wx.MessageDialog(self.main.gui, msg, 'Error', wx.OK | wx.NO_DEFAULT | wx.ICON_ERROR | wx.STAY_ON_TOP)
            dlg.ShowModal()
            raise RuntimeError
        f.fbbAutoConfig(nr_of_birds)
        time.sleep(0.01)
        f.fbbReset()
        time.sleep(0.01)
        f.position_matrix()
        time.sleep(0.01)
        f.run()
        return f

    def ThreadedFOBsetup(self, *args):
        comport = args[0]
        gsi = args[1]
        self.fob_dict[gsi] = self.init_FOB_parallel(comport)
        
    ##  Creates an object for and starts a sensor.
    #   @param[in]  port        The comport where the sensor is plugged into.
    def init_FOB_parallel(self,port):
        f = fob.FOB()
        time.sleep(0.01)
        comport = port-1
        baudrate = 19200#19200
        # Default timeout in seconds.
        timeout = 5
        try: 
            f.open(comport, baudrate, timeout)
            time.sleep(0.01)
            f.fbbAutoConfig(1)
            f.position_matrix()
            time.sleep(0.01)
            f.run()
            return f
        except:
            msg = "Could not open FoB connections: COM port %d does not exist on this computer.\nPlease reconfigure the FobDevice section of the config.ini file." %comport
            print msg
            return -1
            #~ dlg = wx.MessageDialog(self.main.gui, msg, 'Error', wx.OK | wx.NO_DEFAULT | wx.ICON_ERROR | wx.STAY_ON_TOP)
            #~ dlg.ShowModal()
            #~ raise RuntimeError
        

    ##  Puts FoB to sleep. You may only call this method if you inherit from DriverInterface.
    def sleep(self):
        if not isinstance(self, DriverInterface):
            raise TypeError("FoBDriver.sleep requires DriverInterface.")

        for gsi,object in self.flock.items():
            object.sleep()

    ##  Runs XRC. You may only call this method if you inherit from DriverInterface.
    def run(self):
        if not isinstance(self, DriverInterface):
            raise TypeError("FoBDriver.run requires DriverInterface.")
        self.flock[1].run()

    ##  Closes communication. You may only call this method if you inherit from DriverInterface.
    def close(self):
        if not isinstance(self, DriverInterface):
            raise TypeError("FoBDriver.close requires DriverInterface.")
        if self.fobrunning:
            for gsi,object in self.flock.items():
                object.sleep()
                object.close()
            self.fileWriter.close()

    def closeBLFile(self):
        self.fileWriter.close()
        
    ##  Returns raw FoB data of the sensor with the corresponding ID. You may only call this method if you inherit from DriverInterface.
    #   @param[in] id   The sensor ID from which we are requesting data.
    def get_sample(self, id):
        if not isinstance(self, DriverInterface):
            raise TypeError("FoBDriver.get_sample requires DriverInterface.")
        sample = []
        sensorstring = id.split('_')
        # flock is a dict with key gsi, value fob object.
        sample = self.flock[int(sensorstring[1])].point()
        # Write the rawData to a file
        #~ print (sensorstring[1] + '  %.2f  %.2f  %.2f  \r\n   %.2f  %.2f  %.2f  \r\n   %.2f  %.2f  %.2f  \r\n   %.2f  %.2f  %.2f  \r\n\r\n' % (sample[0], sample[1], sample[2], sample[3], sample[4],sample[5], sample[6], sample[7], sample[8], sample[9],sample[10],sample[11]))
        #~ print ('\n')
        try:
            if (self.main.DATA_WRITE or self.main.FAST_ACQUISITION) and self.main.states.visualization.recordingState.RECORDING:
                self.rawoutputlist.append(sensorstring[1] + '  %.2f  %.2f  %.2f  \r\n   %.2f  %.2f  %.2f  \r\n   %.2f  %.2f  %.2f  \r\n   %.2f  %.2f  %.2f  \r\n\r\n' % (sample[0], sample[1], sample[2], sample[3], sample[4],sample[5], sample[6], sample[7], sample[8], sample[9],sample[10],sample[11]))
        except (AttributeError, ValueError):
            pass
            
        return sample

    ##  Calculates the rotations x, y, and z resp. around the x-, y-, and z-axis from matrix R.
    #   You may only call this method if you inherit from DriverInterface.
    def __RotXYZ(self, R):
        if not isinstance(self, DriverInterface):
            raise TypeError("FoBDriver.__init__ requires DriverInterface.")

        y1 = asin(R[0,2])
        sz = -R[0,1] / cos(y1)
        cz =  R[0,0] / cos(y1)
        z1 = atan2(sz,cz)
        sx = -R[1,2] / cos(y1)
        cx =  R[2,2] / cos(y1)
        x1 = atan2(sx,cx)

        if y1 >= 0:
            y2 = pi - y1
        else:
            y2 = -pi -y1

        sz = -R[0,1] / cos(y2)
        cz = R[0,0] / cos(y2)
        z2 = atan2(sz,cz)
        sx = -R[1,2] / cos(y2)
        cx = R[2,2] / cos(y2)
        x2 = atan2(sx,cx)

        if (-pi/2 <= y1) and (y1 <= pi/2):
            return (x1, y1, z1)
        else:
            return (x2, y2, z2)

    ##  Calculates the endpoint of the stylus.
    #   The position and rotation data was not in the correct format and this resulted in a matrix return, this is now fixed.
    #   You may only call this method if you inherit from DriverInterface.
    def __StylusCompensation(self, sensorPosition, sensorRotation):
        if not isinstance(self, DriverInterface):
            raise TypeError("FoBDriver.__init__ requires DriverInterface.")

        return sensorPosition + transpose(sensorRotation) * self.__stylusLength

    ##  Calculates the relative vector of the bony landmark w.r.t. the sensor.
    #   You may only call this method if you inherit from DriverInterface.
    def __CalcBLVec(self, bonyLandmarkPosition, sensorPosition, sensorRotation):
        if not isinstance(self, DriverInterface):
            raise TypeError("FoBDriver.__init__ requires DriverInterface.")
        return sensorRotation * (bonyLandmarkPosition - sensorPosition)

    ##  Calculate position vector of bony landmark relative to its sensors.
    #   You may only call this method if you inherit from DriverInterface.
    def GetBLRelativePosition(self, stylusSensor, sampleSensor):
        if not isinstance(self, DriverInterface):
            raise TypeError("FoBDriver.__init__ requires DriverInterface.")

        relativeSensor = None

        globalBLpos = self.__StylusCompensation(stylusSensor.position,stylusSensor.rotation)

        blvec = self.__CalcBLVec(globalBLpos, sampleSensor.position, sampleSensor.rotation)

        return blvec.mean(axis=1)

    ##  Calculates the global vector of the bony landmark w.r.t. the transmitter coordinate system,
    #   using the local vector of the bony landmark and global vector of the sensor.
    #   You may only call this method if you inherit from DriverInterface.
    def RelativeToGlobal(self, relativePosition, sensorPosition, sensorRotation):
        if not isinstance(self, DriverInterface):
            raise TypeError("FoBDriver.__init__ requires DriverInterface.")

        return sensorPosition + transpose(sensorRotation) * relativePosition

    ##  Gets position and rotation data for all sensors.
    #   You may only call this method if you inherit from DriverInterface.
    def SampleSensors(self):
        if not isinstance(self, DriverInterface):
            raise TypeError("FoBDriver.__init__ requires DriverInterface.")

        #~ start = time.clock()
        for sensor in self.main.sensors:
            thread = threading.Thread(group = None, target = self.ThreadedSampler, args = ([sensor]))
            thread.start()
        while (len(threading.enumerate())>1):
            time.sleep(0.003)
        #~ eind = time.clock()
        #~ print eind-start
        
        #~ for sensor in self.main.sensors:
            #~ ##~ start = time.clock()
            #~ m = self.main.driver.get_sample(sensor.ID)
            #~ ##~ eind = time.clock()
            #~ ##~ print eind-start
            #~ if not self.main.FAST_ACQUISITION:
                #~ sensor.rawPosition = m[0:3]
                #~ sensor.rawRotation = m[3:13]
                #~ sensor = self.CalibrateSingleSensor(sensor)
        
    def ThreadedSampler(self, *args):
        sensor = args[0]
        m = self.main.driver.get_sample(sensor.ID)
        sensor.rawPosition = m[0:3]
        sensor.rawRotation = m[3:13]
        sensor = self.CalibrateSingleSensor(sensor)

    ##  Append a new 'OnPedal transform' to the "bony landmark - sensor mapping": imSensor.
    #   You may only call this method if you inherit from DriverInterface.
    #   @param[in]  blIndex     The bony landmark number.
    #   @param[in]  sensorID    The sensor number.
    def FetchBL_SensorPositionRotation(self, blIndex, sensorID):
        if not isinstance(self, DriverInterface):
            raise TypeError("FoBDriver.__init__ requires DriverInterface.")


        # Get the stylus, now from the self.main.sensors, not form the hardcoded "fob0_2" ID.
        stylusSample = self.get_sample(self.main.sensors[0].ID)
        outputStylusSensor = Sensor(self.main.sensors[0].ID)
        outputStylusSensor.rawPosition = stylusSample[0:3]
        outputStylusSensor.rawRotation = stylusSample[3:13]
        outputStylusSensor = self.CalibrateSingleSensor(outputStylusSensor)

        # Get the measurements for the bony landmark with sensorID id.
        sensorTransform = self.get_sample(sensorID)
        outputSensor = Sensor(sensorID)
        outputSensor.rawPosition = sensorTransform[0:3]
        outputSensor.rawRotation = sensorTransform[3:13]
        outputSensor = self.CalibrateSingleSensor(outputSensor)

        # Outputsensor now has a calibrated position and rotation. Add to imSensors.
        if len(self.imSensors[blIndex]) == 0:
            self.imSensors[blIndex] = [outputSensor]
            self.stylus[blIndex] =  [outputStylusSensor]
        else:
            self.imSensors[blIndex].append(outputSensor)
            self.stylus[blIndex].append(outputStylusSensor)

    ##  This is the new version of CalibrateSensorSSS, because we are only going to calibrate one sensor now.
    #   You may only call this method if you inherit from DriverInterface.
    #   @param[in]  singleSensor    The sensor we are calibrating.
    def CalibrateSingleSensor(self, singleSensor):
        if not isinstance(self, DriverInterface):
            raise TypeError("FoBDriver.__init__ requires DriverInterface.")

        sensorstring = singleSensor.ID.split('_')

        sensornr = int(sensorstring[1])

        # Insert the rawpositions in a matrix.
        rawData = matrix("[%f %f %f]" % (singleSensor.rawPosition[0],singleSensor.rawPosition[1],singleSensor.rawPosition[2]))

        # Include the rawrotations in the same matrix.
        rawData = vstack((rawData, matrix("[%f %f %f;%f %f %f;%f %f %f]" % (singleSensor.rawRotation[0], singleSensor.rawRotation[1], singleSensor.rawRotation[2],
                                                                            singleSensor.rawRotation[3], singleSensor.rawRotation[4], singleSensor.rawRotation[5],
                                                                            singleSensor.rawRotation[6], singleSensor.rawRotation[7], singleSensor.rawRotation[8]))))
        # Calibrate the data.
        newData = self.__CalibrateRawData(rawData)

        # Set the newly caclulated positions and rotation in the sensor.
        pos = transpose(newData[0])
        rot = transpose(newData[1:4])
        singleSensor.position = pos
        singleSensor.rotation = rot

        return singleSensor

    ## Applies the calibration procedure on the raw FoB data using the given calibration file path.
    #
    #  This is called by CalibrateSensors after having built up the rawData structure.
    #  You may only call this method if you inherit from DriverInterface.
    #
    #  @param[in] rawData Matrix with 4 columns of 'raw' fob data (sensor_id, pos, pos, pos) or (sensor_id, rot, rot, rot).
    #
    #  For each sensor id, there is one line with position, and three with the complete rotation matrix.
    #  After a complete sequence of all sensor IDs, the next frame is added.
    def __CalibrateRawData(self, rawData):
          if not isinstance(self, DriverInterface):
              raise TypeError("FoBDriver.__init__ requires DriverInterface.")

          # First iterate through all the POSITION lines. 
          # (the skip=4) in the range skips all the rotation matrices.
          for i in range(0,size(rawData,0),4):
              # Transpose rotation matrix, determine rotations x,y,z around x-, y- and z- axes.
              (x,y,z) = self.__RotXYZ(transpose(rawData[i:i+3,0:3]))
              # Create matrix c [posx, posy, posz, rotx, roty rotz].
              a = matrix("[%f %f %f]" % (x,y,z))
              b = rawData[i,0:3]
              c = concatenate((b,a),1)
              # Append matrix c to reg.
              if i > 0:
                  reg = concatenate((reg,c))
              else:
                  reg = c

          # For EACH row (of 6 elements) do:
          # aa ab ac bb bc cc (done here for 3; just determine all combinations).
          # Store result vector / matrix in T_cross. We could do this per sample.
          for i in range(0,6):
              for j in range(i,6):
                  if (i==0) and (j==0):
                      T_cross = multiply(reg[:,i],reg[:,j])
                  else:
                      T_cross = concatenate((T_cross, multiply(reg[:,i],reg[:,j])),1)
    
    
          # Determine magnitude of position vector.
          # sqrt(x^2 + y^2 + z^2)
          af = sqrt(multiply(reg[:,0],reg[:,0]) + multiply(reg[:,1],reg[:,1]) + multiply(reg[:,2],reg[:,2]))
    
          # T has 29 columns.
          #
          # 1 x y z magpos rotx roty rotz aa bb cc bb bc cc
          # 1 x y z magpos rotx roty rotz ...
          # 1 x y z magpos rotx roty rotz ...
          T = hstack((ones((size(reg,0),1)), reg[:,0:3], af, reg[:,3:6], T_cross))

          # th_nonline has 29 rows and 3 columns.
          th_nonlin = self.main.calibrationMatrix
          # error has size(reg,0) rows and 3 columns.
          error = T * th_nonlin

          # Copy complete rawData (only positions have been converted to millimetres).
          newData = rawData.copy()
          # For each position (note the skip=4), add x, y and z error to the position, divide the whole thing by 10 to get centimetres.
          for i in range(0, size(rawData,0), 4):
              newData[i,0:3] = (rawData[i,0:3] +error[(i)/4,:])/10

          return newData